#include "Graph.hpp"
#include <map>
#include <set>
#include <limits>
#include <queue>
#include <tuple>

template <typename T>
auto createGraph()
{
    Graph<T> G(9);
    std::map<unsigned, std::vector<std::pair<size_t, T>>> edges;
    edges[1] = { {2, 2}, {5, 3} };
    edges[2] = { {1, 2}, {5, 5}, {4, 1} };
    edges[3] = { {4, 2}, {7, 3} };
    edges[4] = { {2, 1}, {3, 2}, {5, 2}, {6, 4}, {8, 5} };
    edges[5] = { {1, 3}, {2, 5}, {4, 2}, {8, 3} };
    edges[6] = { {4, 4}, {7, 4}, {8, 1} };
    edges[7] = { {3, 3}, {6, 4} };
    edges[8] = { {4, 5}, {5, 3}, {6, 1} };
    for (auto& i : edges)
        for (auto& j : i.second)
            G.addEdge(Edge<T>{ i.first, j.first, j.second });
    return G;
}

template <typename T>
struct Label {
    size_t vertexID;
    T distance;
    Label(size_t id, T d): vertexID(id), distance(d) {}
    inline bool operator<(const Label<T>& that) const {
        return distance < that.distance;
    }
    inline bool operator>(const Label<T>& that) const {
        return distance > that.distance;
    }
    inline bool operator()(const Label<T>& that) const {
        return operator>(that);
    }
};

template <typename T>
auto primMST(const Graph<T>& G, size_t src) {
    std::priority_queue<
        Label<T>, 
        vector<Label<T>>, 
        std::greater<Label<T>>
    > heap; //使用了小根堆，因此每次优先弹出距离父节点最近的顶点
    std::set<size_t> visited;
    vector<T> distance(G.vertices(), std::numeric_limits<T>::max());
    std::vector<size_t> MST;
    heap.emplace(src, 0); //起点的路径长度为0
    T totalWeight{};

    while (!heap.empty()) {
        auto currentVertex = heap.top();
        heap.pop();
        //如果未被访问过，则访问
        if (visited.find(currentVertex.vertexID) == visited.end()) {
            //将当前顶点存入最小生成树中
            MST.push_back(currentVertex.vertexID);
            for (const auto& e : G.edgesFromVertex(currentVertex.vertexID)) {
                auto neighborID = e.dest;
                auto newDistance = e.weight;
                //如果当前边终点到当前顶点的距离小于终点到
                if (newDistance < distance[neighborID]) {
                    heap.emplace(neighborID, newDistance);
                    distance[e.dest] = newDistance;
                    totalWeight += newDistance;
                }
            }
        }

        visited.insert(currentVertex.vertexID);
    }
    return std::tuple<vector<size_t>, T>(MST, totalWeight);
}

template <typename T>
void testPrimMST() {
    auto G = createGraph<T>();
    cout << G << endl;
    auto tupleMstWeight = primMST<T>(G, 1);
    auto MST = std::get<0>(tupleMstWeight);
    auto totalWeight = std::get<1>(tupleMstWeight);
    cout << "Minimum Spanning Tree:\n";
    for (const auto& v : MST)
        cout << v << endl;
    cout << "Total weight = " << totalWeight << endl;
}

int main(int argc, const char** argv) {
    using T = uint32_t;
    testPrimMST<T>();
    return 0;
}